import asyncio
import aiohttp
import os
import tempfile
import shutil
import re
import logging
from typing import List, Dict, Optional, Tuple
from urllib.parse import urljoin, urlparse
import subprocess
from pathlib import Path

logger = logging.getLogger(__name__)

class HLSDownloader:
    """HLS (M3U8) 视频下载和合成工具类"""
    
    def __init__(self, temp_dir: Optional[str] = None, cleanup_temp: bool = True):
        """
        初始化HLS下载器
        
        Args:
            temp_dir: 临时文件目录，如果为None则使用系统临时目录
            cleanup_temp: 是否在完成后清理临时文件
        """
        self.temp_dir = temp_dir or tempfile.gettempdir()
        self.cleanup_temp = cleanup_temp
        self.session = None
        
    async def __aenter__(self):
        """异步上下文管理器入口"""
        self.session = aiohttp.ClientSession(
            timeout=aiohttp.ClientTimeout(total=60, connect=15, sock_read=30),
            headers={
                'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
                'Accept': '*/*',
                'Accept-Language': 'en-US,en;q=0.9',
                'Accept-Encoding': 'gzip, deflate, br',
                'Connection': 'keep-alive',
                'Sec-Fetch-Dest': 'empty',
                'Sec-Fetch-Mode': 'cors',
                'Sec-Fetch-Site': 'cross-site'
            },
            connector=aiohttp.TCPConnector(
                limit=100,
                limit_per_host=10,
                ttl_dns_cache=300,
                use_dns_cache=True,
                ssl=False  # 对于Google Video可能需要禁用SSL验证
            )
        )
        return self
        
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器退出"""
        if self.session:
            await self.session.close()
            
    def _parse_m3u8(self, content: str, base_url: str) -> List[str]:
        """
        解析M3U8播放列表内容
        
        Args:
            content: M3U8文件内容
            base_url: 基础URL，用于构建完整的片段URL
            
        Returns:
            片段URL列表
        """
        segments = []
        lines = content.strip().split('\n')
        
        for line in lines:
            line = line.strip()
            # 跳过注释和空行
            if line.startswith('#') or not line:
                continue
                
            # 如果是相对URL，构建完整URL
            if line.startswith('http'):
                segments.append(line)
            else:
                # 处理相对路径
                full_url = urljoin(base_url, line)
                segments.append(full_url)
                
        logger.info(f"解析M3U8播放列表完成，找到 {len(segments)} 个视频片段")
        return segments
        
    async def _download_playlist(self, m3u8_url: str, max_retries: int = 3) -> Tuple[List[str], str]:
        """
        下载M3U8播放列表文件，支持重试机制
        
        Args:
            m3u8_url: M3U8播放列表URL
            max_retries: 最大重试次数
            
        Returns:
            (片段URL列表, 基础URL)
        """
        last_error = None
        
        for attempt in range(max_retries + 1):
            try:
                logger.info(f"尝试下载M3U8播放列表 (第{attempt + 1}次): {m3u8_url[:100]}...")
                
                async with self.session.get(m3u8_url) as response:
                    logger.info(f"HTTP响应状态: {response.status}, Content-Type: {response.headers.get('Content-Type', 'Unknown')}")
                    
                    if response.status != 200:
                        error_msg = f"HTTP错误，状态码: {response.status}"
                        if response.status == 403:
                            error_msg += " (访问被拒绝，可能是权限或地区限制)"
                        elif response.status == 404:
                            error_msg += " (资源不存在，链接可能已过期)"
                        elif response.status >= 500:
                            error_msg += " (服务器错误，可能是临时问题)"
                        raise Exception(error_msg)
                        
                    content = await response.text()
                    logger.info(f"成功下载M3U8播放列表，大小: {len(content)} 字符")
                    
                    # 检查内容是否为有效的M3U8格式
                    if not content.strip().startswith('#EXTM3U'):
                        logger.warning(f"M3U8内容预览: {content[:200]}...")
                        raise Exception("响应内容不是有效的M3U8格式")
                    
                    # 获取基础URL
                    parsed_url = urlparse(m3u8_url)
                    base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
                    if parsed_url.path:
                        # 获取目录路径
                        base_path = '/'.join(parsed_url.path.split('/')[:-1])
                        if base_path:
                            base_url += base_path + '/'
                        else:
                            base_url += '/'
                    
                    # 解析片段URL
                    segments = self._parse_m3u8(content, base_url)
                    
                    return segments, base_url
                    
            except asyncio.TimeoutError as e:
                last_error = f"下载超时: {str(e)}"
                logger.warning(f"第{attempt + 1}次尝试超时")
                
            except aiohttp.ClientError as e:
                last_error = f"网络连接错误: {str(e)}"
                logger.warning(f"第{attempt + 1}次尝试网络错误: {str(e)}")
                
            except Exception as e:
                last_error = f"下载失败: {str(e)}"
                logger.warning(f"第{attempt + 1}次尝试失败: {str(e)}")
            
            # 如果不是最后一次尝试，等待后重试
            if attempt < max_retries:
                wait_time = 2 ** attempt  # 指数退避
                logger.info(f"等待 {wait_time} 秒后重试...")
                await asyncio.sleep(wait_time)
        
        # 所有重试都失败了
        error_msg = f"经过 {max_retries + 1} 次尝试后仍然失败，最后错误: {last_error}"
        logger.error(error_msg)
        raise Exception(error_msg)
            
    async def _download_segment(self, segment_url: str, output_path: str, segment_index: int, max_retries: int = 2) -> bool:
        """
        下载单个视频片段，支持重试机制
        
        Args:
            segment_url: 片段URL
            output_path: 输出文件路径
            segment_index: 片段索引
            max_retries: 最大重试次数
            
        Returns:
            是否下载成功
        """
        for attempt in range(max_retries + 1):
            try:
                async with self.session.get(segment_url) as response:
                    if response.status != 200:
                        logger.warning(f"片段 {segment_index} 第{attempt + 1}次尝试，状态码: {response.status}")
                        if attempt == max_retries:
                            logger.error(f"下载片段 {segment_index} 失败，最终状态码: {response.status}")
                            return False
                        continue
                        
                    # 保存片段文件
                    with open(output_path, 'wb') as f:
                        async for chunk in response.content.iter_chunked(8192):
                            f.write(chunk)
                            
                    file_size = os.path.getsize(output_path)
                    logger.debug(f"片段 {segment_index} 下载完成，大小: {file_size} 字节")
                    return True
                    
            except Exception as e:
                logger.warning(f"片段 {segment_index} 第{attempt + 1}次尝试失败: {str(e)}")
                if attempt == max_retries:
                    logger.error(f"下载片段 {segment_index} 最终失败: {str(e)}")
                    return False
                    
                # 短暂等待后重试
                await asyncio.sleep(0.5)
                
        return False
            
    async def _download_all_segments(self, segments: List[str], temp_folder: str) -> List[str]:
        """
        并发下载所有视频片段
        
        Args:
            segments: 片段URL列表
            temp_folder: 临时文件夹路径
            
        Returns:
            成功下载的片段文件路径列表
        """
        logger.info(f"开始下载 {len(segments)} 个视频片段...")
        
        # 创建下载任务
        tasks = []
        segment_files = []
        
        for i, segment_url in enumerate(segments):
            segment_file = os.path.join(temp_folder, f"segment_{i:04d}.ts")
            segment_files.append(segment_file)
            
            task = self._download_segment(segment_url, segment_file, i)
            tasks.append(task)
            
        # 并发执行下载任务（限制并发数）
        semaphore = asyncio.Semaphore(5)  # 最多5个并发下载
        
        async def download_with_semaphore(task):
            async with semaphore:
                return await task
                
        results = await asyncio.gather(*[download_with_semaphore(task) for task in tasks])
        
        # 检查下载结果
        successful_files = []
        failed_count = 0
        
        for i, (success, file_path) in enumerate(zip(results, segment_files)):
            if success and os.path.exists(file_path):
                successful_files.append(file_path)
            else:
                failed_count += 1
                logger.error(f"片段 {i} 下载失败")
                
        if failed_count > 0:
            logger.warning(f"有 {failed_count} 个片段下载失败")
            
        logger.info(f"片段下载完成，成功: {len(successful_files)}, 失败: {failed_count}")
        return successful_files
        
    def _merge_segments_with_ffmpeg(self, segment_files: List[str], output_path: str) -> bool:
        """
        使用FFmpeg合并视频片段
        
        Args:
            segment_files: 片段文件路径列表
            output_path: 输出视频文件路径
            
        Returns:
            是否合并成功
        """
        try:
            logger.info(f"使用FFmpeg合并 {len(segment_files)} 个片段...")
            
            # 创建文件列表
            temp_list_file = os.path.join(os.path.dirname(output_path), 'segments_list.txt')
            
            with open(temp_list_file, 'w', encoding='utf-8') as f:
                for segment_file in segment_files:
                    # FFmpeg需要转义路径中的特殊字符
                    escaped_path = segment_file.replace('\\', '\\\\').replace("'", "\\'")
                    f.write(f"file '{escaped_path}'\n")
            
            # 构建FFmpeg命令
            cmd = [
                'ffmpeg',
                '-f', 'concat',
                '-safe', '0',
                '-i', temp_list_file,
                '-c', 'copy',  # 复制流，不重新编码
                '-y',  # 覆盖输出文件
                output_path
            ]
            
            # 执行FFmpeg命令
            logger.debug(f"执行FFmpeg命令: {' '.join(cmd)}")
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            # 清理临时文件列表
            if os.path.exists(temp_list_file):
                os.remove(temp_list_file)
                
            if result.returncode == 0:
                output_size = os.path.getsize(output_path)
                logger.info(f"FFmpeg合并成功，输出文件大小: {output_size} 字节")
                return True
            else:
                logger.error(f"FFmpeg合并失败: {result.stderr}")
                return False
                
        except Exception as e:
            logger.error(f"FFmpeg合并过程发生错误: {str(e)}")
            return False
            
    def _merge_segments_direct(self, segment_files: List[str], output_path: str) -> bool:
        """
        直接合并TS片段（无需重新编码）
        
        Args:
            segment_files: 片段文件路径列表  
            output_path: 输出视频文件路径
            
        Returns:
            是否合并成功
        """
        try:
            logger.info(f"直接合并 {len(segment_files)} 个TS片段...")
            
            with open(output_path, 'wb') as output_file:
                for i, segment_file in enumerate(segment_files):
                    if os.path.exists(segment_file):
                        with open(segment_file, 'rb') as segment:
                            shutil.copyfileobj(segment, output_file)
                        logger.debug(f"合并片段 {i}")
                    else:
                        logger.warning(f"片段文件不存在: {segment_file}")
                        
            output_size = os.path.getsize(output_path)
            logger.info(f"直接合并成功，输出文件大小: {output_size} 字节")
            return True
            
        except Exception as e:
            logger.error(f"直接合并过程发生错误: {str(e)}")
            return False
            
    async def diagnose_hls_url(self, hls_url: str) -> Dict[str, any]:
        """
        诊断HLS链接的可访问性和基本信息
        
        Args:
            hls_url: HLS链接
            
        Returns:
            诊断结果
        """
        try:
            logger.info(f"开始诊断HLS链接: {hls_url[:100]}...")
            
            # 基本URL格式检查
            parsed_url = urlparse(hls_url)
            url_info = {
                "scheme": parsed_url.scheme,
                "netloc": parsed_url.netloc,
                "path": parsed_url.path,
                "is_https": parsed_url.scheme == "https",
                "domain": parsed_url.netloc.split(':')[0]
            }
            
            # 尝试HEAD请求检查可访问性
            try:
                async with self.session.head(hls_url) as response:
                    head_info = {
                        "status": response.status,
                        "headers": dict(response.headers),
                        "accessible": response.status < 400
                    }
            except Exception as e:
                head_info = {
                    "status": None,
                    "error": str(e),
                    "accessible": False
                }
            
            # 尝试下载少量内容
            content_info = {}
            try:
                async with self.session.get(hls_url) as response:
                    if response.status == 200:
                        content = await response.text()
                        content_info = {
                            "status": response.status,
                            "content_length": len(content),
                            "content_type": response.headers.get('Content-Type', 'Unknown'),
                            "is_m3u8": content.strip().startswith('#EXTM3U'),
                            "content_preview": content[:300] if content else None
                        }
                    else:
                        content_info = {
                            "status": response.status,
                            "error": f"HTTP {response.status}",
                            "accessible": False
                        }
            except Exception as e:
                content_info = {
                    "error": str(e),
                    "accessible": False
                }
            
            return {
                "success": True,
                "url_info": url_info,
                "head_request": head_info,
                "content_info": content_info,
                "recommendations": self._get_recommendations(url_info, head_info, content_info)
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }
    
    def _get_recommendations(self, url_info: dict, head_info: dict, content_info: dict) -> List[str]:
        """根据诊断结果生成建议"""
        recommendations = []
        
        if not url_info.get("is_https"):
            recommendations.append("建议使用HTTPS链接以提高安全性")
            
        if not head_info.get("accessible", True):
            recommendations.append("链接无法访问，请检查网络连接或链接有效性")
            
        if content_info.get("status") == 403:
            recommendations.append("访问被拒绝，可能需要特定的请求头或身份验证")
        elif content_info.get("status") == 404:
            recommendations.append("资源不存在，链接可能已过期")
        elif content_info.get("status", 0) >= 500:
            recommendations.append("服务器错误，请稍后重试")
            
        if not content_info.get("is_m3u8", False):
            recommendations.append("响应内容不是有效的M3U8格式")
            
        if "googlevideo.com" in url_info.get("domain", ""):
            recommendations.append("Google Video链接可能有地区或时间限制")
            
        return recommendations
            
    async def download_hls_video(
        self, 
        m3u8_url: str, 
        output_path: str, 
        use_ffmpeg: bool = True
    ) -> Dict[str, any]:
        """
        下载HLS视频并合成完整视频文件
        
        Args:
            m3u8_url: M3U8播放列表URL
            output_path: 输出视频文件路径
            use_ffmpeg: 是否使用FFmpeg合并（推荐），False则使用直接合并
            
        Returns:
            下载结果字典
        """
        temp_folder = None
        
        try:
            # 创建临时文件夹
            temp_folder = tempfile.mkdtemp(prefix='hls_download_', dir=self.temp_dir)
            logger.info(f"创建临时文件夹: {temp_folder}")
            
            # 下载并解析M3U8播放列表
            segments, base_url = await self._download_playlist(m3u8_url)
            
            if not segments:
                return {
                    "success": False,
                    "error": "M3U8播放列表为空或解析失败"
                }
                
            # 下载所有片段
            segment_files = await self._download_all_segments(segments, temp_folder)
            
            if not segment_files:
                return {
                    "success": False,
                    "error": "所有视频片段下载失败"
                }
                
            # 确保输出目录存在
            output_dir = os.path.dirname(output_path)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir, exist_ok=True)
                
            # 合并片段
            if use_ffmpeg:
                merge_success = self._merge_segments_with_ffmpeg(segment_files, output_path)
            else:
                merge_success = self._merge_segments_direct(segment_files, output_path)
                
            if not merge_success:
                return {
                    "success": False,
                    "error": "视频片段合并失败"
                }
                
            # 验证输出文件
            if not os.path.exists(output_path):
                return {
                    "success": False,
                    "error": "输出视频文件不存在"
                }
                
            file_size = os.path.getsize(output_path)
            
            return {
                "success": True,
                "output_path": output_path,
                "file_size": file_size,
                "segments_count": len(segments),
                "downloaded_segments": len(segment_files),
                "temp_folder": temp_folder
            }
            
        except Exception as e:
            logger.error(f"HLS视频下载失败: {str(e)}")
            return {
                "success": False,
                "error": f"下载过程发生错误: {str(e)}"
            }
            
        finally:
            # 清理临时文件
            if self.cleanup_temp and temp_folder and os.path.exists(temp_folder):
                try:
                    shutil.rmtree(temp_folder)
                    logger.info(f"清理临时文件夹: {temp_folder}")
                except Exception as e:
                    logger.warning(f"清理临时文件失败: {str(e)}")
                    
    def check_ffmpeg_available(self) -> bool:
        """
        检查FFmpeg是否可用
        
        Returns:
            FFmpeg是否可用
        """
        try:
            result = subprocess.run(['ffmpeg', '-version'], 
                                  capture_output=True, text=True)
            return result.returncode == 0
        except FileNotFoundError:
            return False
            
    @staticmethod
    def is_hls_url(url: str) -> bool:
        """
        检查URL是否为HLS格式
        
        Args:
            url: 要检查的URL
            
        Returns:
            是否为HLS URL
        """
        return 'manifest' in url and ('m3u8' in url or 'hls' in url)


# 便捷函数
async def download_hls_video(
    m3u8_url: str, 
    output_path: str, 
    use_ffmpeg: bool = True,
    cleanup_temp: bool = True
) -> Dict[str, any]:
    """
    便捷函数：下载HLS视频
    
    Args:
        m3u8_url: M3U8播放列表URL
        output_path: 输出视频文件路径
        use_ffmpeg: 是否使用FFmpeg合并
        cleanup_temp: 是否清理临时文件
        
    Returns:
        下载结果字典
    """
    async with HLSDownloader(cleanup_temp=cleanup_temp) as downloader:
        return await downloader.download_hls_video(m3u8_url, output_path, use_ffmpeg) 